import openke
from openke.config import Trainer, Tester
from openke.module.model import TransE
from openke.module.loss import MarginLoss
from openke.module.strategy import NegativeSampling
from openke.data import TrainDataLoader, TestDataLoader ,PyTorchTrainDataLoader

if __name__ == '__main__':
	# dataloader for training
	train_dataloader = PyTorchTrainDataLoader(
		in_path = "../benchmarks/MY/医学知识/",
		nbatches = 10,
		threads = 0,
		sampling_mode = "normal",
		bern_flag = 1,
		filter_flag = 1,
		neg_ent = 25,
		neg_rel = 0)

	# dataloader for test
	# test_dataloader = TestDataLoader("./benchmarks/FB15K237/", "link")

	# define the model
	transe = TransE(
		ent_tot = train_dataloader.get_ent_tot(),
		rel_tot = train_dataloader.get_rel_tot(),
		dim = 200,
		p_norm = 1,
		norm_flag = True)


	# define the loss function
	model = NegativeSampling(
		model = transe,
		loss = MarginLoss(margin = 5.0),
		batch_size = train_dataloader.get_batch_size()
	)

	# train the model
	trainer = Trainer(model = model, data_loader = train_dataloader, train_times = 400, alpha = 1.0, use_gpu = True)
	trainer.run()
	transe.save_checkpoint('../checkpoint/医学知识/transe.ckpt')

	'''
	# test the model
	transe.load_checkpoint('./checkpoint/医学知识/transe.ckpt')
	tester = Tester(model = transe, data_loader = test_dataloader, use_gpu = True)
	tester.run_link_prediction(type_constrain = False)
	'''